/*
 * Copyright (c) 2022 HiSilicon (Shanghai) Technologies CO., LIMITED.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * 该cpp文件基于OpenCV实现了网球检测功能。为了保证FPS的帧数，
 * 我们设计的原则是IVE（Intelligent Video Engine）+AI CPU结合使用，即IVE不支持的算子
 * 通过AI CPU进行计算，否则走IVE硬件加速模块进行处理。并将检测的结果通过VGS标记出来。
 *
 * The cpp file implements the tennis ball detection function based on OpenCV.
 * In order to ensure the number of FPS frames,
 * The principle of our design is the combination of IVE (Intelligent Video Engine) + AI CPU,
 * that is, operators not supported by IVE are calculated by AI CPU, otherwise,
 * the IVE hardware acceleration module is used for processing.
 * And the detection results are marked by VGS.
 */

#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include "kalman_tracking.hpp"

#include "sample_comm_ive.h"

using namespace std;
using namespace cv;

#define max(a, b)           (((a) > (b)) ? (a) : (b))
#define min(a, b)           (((a) < (b)) ? (a) : (b))

// /* create class */
// KalmanTracking tracker;

/* function : x, y, x, y to x, y, w, h */
static Rect xyxy_to_xywh(RectBox bBox)
{
    Rect roi;
    roi.x = (bBox.xmin + bBox.xmax) / 2.;
    roi.y = (bBox.ymin + bBox.ymax) / 2.;
    roi.width = (bBox.xmax - bBox.xmin);
    roi.height = (bBox.ymax - bBox.ymin);
    return roi;
}

/* function : x, y, w, h to x, y, x, y */
static RectBox xywh_to_xyxy(Rect roi)
{
    RectBox bBox;
    bBox.xmin = roi.x - roi.width / 2.;
    bBox.xmax = roi.x + roi.width / 2.;
    bBox.ymin = roi.y - roi.height / 2.;
    bBox.ymax = roi.y + roi.height / 2.;
    return bBox;
}

/* function : caculate the iou between two rect */
static HI_FLOAT cal_iou(RectBox box1, RectBox box2)
{
    HI_FLOAT iou1, iou2;

    /* Caculate two iou area */
    iou1 = (box1.xmax - box1.xmin + 1.) * (box1.ymax - box1.ymin + 1.);
    iou2 = (box2.xmax - box2.xmin + 1.) * (box2.ymax - box2.ymin + 1.);

    /* Caculate the interiou */
    HI_FLOAT interIou;

    HI_FLOAT xmin = max(box1.xmin, box2.xmin);
    HI_FLOAT ymin = max(box1.ymin, box2.ymin);
    HI_FLOAT xmax = min(box1.xmax, box2.xmax);
    HI_FLOAT ymax = min(box1.ymax, box2.ymax);

    HI_FLOAT inter_h = max(ymax - ymin + 1., 0.);
    HI_FLOAT inter_w = max(xmax - xmin + 1., 0.);

    interIou = inter_h * inter_w;

    /* Caculate the iou ratio */
    return (HI_FLOAT)(interIou / (iou1 + iou2 - interIou));
}

/* function : mat to rect */
static Rect mat_to_xywh(Mat matrix)
{
    Rect roi;
    roi.x       = matrix.at<float>(0, 0);
    roi.y       = matrix.at<float>(1, 0);
    roi.width   = matrix.at<float>(2, 0);
    roi.height  = matrix.at<float>(3, 0);
    return roi;
}

/* function : use max iou match target, return measure state matrix */
void KalmanTracking::max_iou_match(SVP_NNIE_YOLOV5_BBOX_S *nodes, HI_BOOL *is_match, HI_FLOAT match_threshold,
                Mat s_statePost, Mat &s_measurement)
{
    HI_FLOAT thresh = match_threshold;
    RectBox matchBox, targetBox, centerBox;
    HI_FLOAT match_iou;
    int calDist, dist = INT_MAX;
    targetBox = xywh_to_xyxy(mat_to_xywh(s_statePost));
    while (nodes != NULL) {
        /* case 1: get match iou */
        match_iou = cal_iou(nodes->bBox, targetBox);
        if (match_iou > thresh) {
            thresh = match_iou;
            matchBox = nodes->bBox;
            *is_match = HI_TRUE;
        }

        /* case 2: get the center target */
        calDist = CaculateDiffCenter(nodes->bBox, 160.0);
        if (calDist < dist) {
            dist = calDist;
            centerBox = nodes->bBox;
        }

        nodes = nodes->next;
    }

    if (*is_match == HI_TRUE) {
        cv::Rect matchRect = xyxy_to_xywh(matchBox);
        HI_FLOAT dx = matchRect.x - xyxy_to_xywh(targetBox).x;
        HI_FLOAT dy = matchRect.y - xyxy_to_xywh(targetBox).y;
        s_measurement = (Mat_<float>(6, 1, CV_32F) << 
            matchRect.x, matchRect.y, matchRect.width, matchRect.height, dx, dy);
    } else {
        cv::Rect centerRect = xyxy_to_xywh(centerBox);
        HI_FLOAT dx = centerRect.x - xyxy_to_xywh(targetBox).x;
        HI_FLOAT dy = centerRect.y - xyxy_to_xywh(targetBox).y;
        if (abs(dx) < 45)
        {
            s_measurement = (Mat_<float>(6, 1, CV_32F) << 
                centerRect.x, centerRect.y, centerRect.width, centerRect.height, dx, dy);
        }
    }
}

void KalmanTracking::init(RectBox targetBox)
{
    /* 状态转移矩阵 A */
    transitionMatrix = (Mat_<float>(6, 6, CV_32F) << 
        1, 0, 0, 0, 1, 0,
        0, 1, 0, 0, 0, 1,
        0, 0, 1, 0, 0, 0,
        0, 0, 0, 1, 0, 0,
        0, 0, 0, 0, 1, 0,
        0, 0, 0, 0, 0, 1);

    /* 状态观测矩阵 H */
    measurementMatrix   = Mat::eye(6, 6, CV_32F);

    /* 过程噪声协方差矩阵 Q */
    processNoiseCov     = Mat::eye(6, 6, CV_32F) * 1.0;

    /* 观测噪声协方差矩阵 R */
    measurementNoiseCov = Mat::eye(6, 6, CV_32F) * 0.05;

    /* 状态估计协方差矩阵 P */
    errorCovPre       = Mat::eye(6, 6, CV_32F);

    /* 初始状态矩阵赋值 [中心x, 中心y, 宽w, 高h, dx, dy].T */
    Rect initRoi = xyxy_to_xywh(targetBox);
    statePre = (Mat_<float>(6, 1, CV_32F) << 
        initRoi.x, initRoi.y, initRoi.width, initRoi.height, 0.0, 0.0);

    /* 后验状态矩阵 */
    statePre.copyTo(statePost);

    /* 后验状态估计 */
    errorCovPre.copyTo(errorCovPost);

    /* 测量值 */
    statePre.copyTo(measurement);

    /* 卡尔曼增益 */
    gain = Mat::zeros(6, 6, CV_32F);

    /* 过程计算值 */
    temp1.create(6, 6, CV_32F);
    temp2.create(6, 6, CV_32F);
    temp3.create(6, 6, CV_32F);
    temp4.create(6, 6, CV_32F);
    temp5.create(6, 1, CV_32F);
    temp6.create(6, 6, CV_32F);
}

RectBox KalmanTracking::KalmanTrackingTarget(SAMPLE_IVE_KCF_S *stKcfInfo, SVP_NNIE_YOLOV5_BBOX_S *nodes, HI_FLOAT match_thresh, HI_BOOL *isMatch)
{
    HI_BOOL max_iou_matched = HI_FALSE;
    RectBox result = stKcfInfo->trackRoi;
    /* confirm the track target, init kalman tracking */
    if (stKcfInfo->enKcfProcStat == KCF_PROC_START) {
        init(stKcfInfo->trackRoi);
        stKcfInfo->enKcfProcStat = KCF_PROC_PROC;
    }
    if (stKcfInfo->enKcfProcStat == KCF_PROC_PROC) {
        max_iou_match(nodes, &max_iou_matched, match_thresh, statePost, measurement);
        if (max_iou_matched == HI_TRUE) {
            /* 先验估计 x'(k) = A*x(k) */
            statePre = transitionMatrix * statePost;

            /* 更新状态估计协方差矩阵P temp1 = A*P(k) */
            temp1 = transitionMatrix * errorCovPost;

            /* P'(k) = temp1*At + Q */
            gemm(temp1, transitionMatrix, 1.0, processNoiseCov, 1.0, errorCovPre, GEMM_2_T);

            /* 计算Kalman gain */
            // temp2 = H*P'(k) mp * dp * dp * dp
            temp2 = measurementMatrix * errorCovPre;
            
            // temp3 = temp2*Ht + R mp * dp * dp * mp  + mp * mp
            gemm(temp2, measurementMatrix, 1.0, measurementNoiseCov, 1.0, temp3, GEMM_2_T);

            // temp4 = inv(temp3)*temp2 = Kt(k) mp * mp * mp *dp
            solve(temp3, temp2, temp4, DECOMP_SVD);

            // K(k) dp * mp
            gain = temp4.t();

            /* 后验估计 */
            // temp5 = z(k) - H*x'(k)  mp * 1 - mp * dp * dp * 1
            temp5 = measurement - measurementMatrix * statePre;

            // x(k) = x'(k) + K(k)*temp5 dp * + dp * mp * mp * 1
            statePost = statePre + gain * temp5;

            /* 更新状态估计协方差矩阵P */
            // P(k) = P'(k) - K(k)*temp2 6 * 6 - dp * mp * mp * dp
            temp6 = Mat::eye(6, 6, CV_32F) - gain * measurementMatrix;
            errorCovPost = temp6 * errorCovPre;
            // errorCovPost = errorCovPre - gain * temp2;
        } else {
            /* 如果IOU匹配失败，此时失去观测值，那么直接使用上一次的最优估计作为先验估计 */
            // statePost = transitionMatrix * statePost;
            /* 如果IOU匹配失败，直接取最近人脸为跟踪目标 */
            statePost = measurement;
        }
        
        /* output result */
        result = xywh_to_xyxy(mat_to_xywh(statePost));
        result.xmin = min(max(result.xmin, 0), YOLO_MODEL_WIDTH);
        result.ymin = min(max(result.ymin, 0), YOLO_MODEL_HEIGHT);
        result.xmax = max(min(result.xmax, YOLO_MODEL_WIDTH), 0);
        result.ymax = max(min(result.ymax, YOLO_MODEL_HEIGHT), 0);
    }
    *isMatch = max_iou_matched;
    return result;
}
